#%%
from typing import Union

import numpy as np
from itertools import *
from functools import *
import operator
import joblib
import datetime
from scipy.stats import norm
import warnings

warnings.filterwarnings("error")
#%%
# Signal Function
def I(x: np.ndarray):
    t = np.zeros([x.shape[0], 2])
    t[:, 0] = x
    return np.max(t, axis=1)

# Expectation Function
def E(x: Union[np.ndarray, list], p: np.ndarray):
    return np.dot(x, p)

# PDF function
def p_norm(x: np.ndarray, mu, sigma):
    init_arr = norm.cdf(x, mu, sigma)
    res_arr = np.zeros(x.shape[0])
    res_arr[0] = init_arr[0]
    res_arr[1:] = np.diff(init_arr)
    res_arr[-1] = 1 - init_arr[-1]
    return res_arr

# Variable
T = 3
alpha = 0.95
# Kt
mu_K_t = 6  # 均值
cv_K_t = 0.3  # 变异系数
sigma_K_t = cv_K_t * mu_K_t  # 标准差
K_t = np.array([4, 5, 6, 7, 8])
p_K_t = p_norm(K_t, mu_K_t, sigma_K_t)
# Dt
mu_t = 5
sigma_t = 0.5
D_t = np.array(list(range(0, 10)))
p_D_t = p_norm(D_t, mu_t, sigma_t)
# Joint Distribution of Kt, Dt
K_D_t = [(K_t_i, D_t_i) for K_t_i, D_t_i in product(K_t, D_t)]
# Consider as independent
p_K_D_t = [p_K_t_i * p_D_t_i for p_K_t_i, p_D_t_i in product(p_K_t, p_D_t)]

p = 19
b = 19
c_e = 10
c_m = 2
h_e = 1
h_m = 0.2
theta = h_m / (1 - alpha)
#%%
def g_1t(x, epoch):
    val = h_e * E(I(x - D_t), p_D_t) + b * E(I(D_t - x), p_D_t) + (c_e - h_m) * x
    if epoch != max_epoch:
        # last cache need to be added for f_0t
        res = alpha * E([f_0t(x - D_t_i, epoch + 1) for D_t_i in D_t], p_D_t)
        if val > np.inf and res < -np.inf:
            val = np.inf
        elif val < -np.inf and res > np.inf:
            val = -np.inf
        else:
            val += res
    return val

def g_2t(x, epoch):
    val = 0
    if epoch != max_epoch:
        # last cache need to be added for f_1t
        val += alpha * E([f_1t(x - D_t_i, epoch + 1) for D_t_i in D_t], p_D_t)
    return val

def g_3t(x, epoch):
    exp = [np.min([np.max([S_t(epoch), x]), np.max([x + K_t_i])])  for K_t_i in K_t]
    val = (h_m - c_m) * x + c_m * E(exp, p_K_t)
    if epoch != max_epoch:
        val += (alpha * E([f_2t(x - D_t_i, epoch + 1) for D_t_i in D_t], p_D_t) +
                alpha * E([f_3t(np.min([np.max([S_t(epoch), x]), x + K_t_i]) - D_t_i, epoch + 1) for K_t_i, D_t_i in K_D_t], p_K_D_t))
    return val


def f_0t(v_0, epoch):
    return - c_e * v_0

def f_1t(v_1, epoch):
    return g_1t(np.max([v_1, S_2t(epoch)]), epoch)

def f_2t(v_2, epoch):
    return g_1t(np.max([v_2, S_2t(epoch)]), epoch) + g_2t(np.max([v_2, S_3t(epoch)]), epoch) - g_1t(np.max([v_2, S_1t(epoch)]), epoch)

def f_3t(v_3, epoch):
    return g_1t(np.max([v_3, S_2t(epoch)]), epoch) + g_2t(np.max([v_3, S_3t(epoch)]), epoch) + g_3t(np.max([v_3, S_4t(epoch)]), epoch) \
           - g_1t(np.max([v_3, S_1t(epoch)]), epoch) - g_2t(np.max([v_3, S_2t(epoch)]), epoch)


# S_t can be cached
def S_t(epoch):
    if epoch != max_epoch:
        return z[np.argmin([c_m * _z] for _z in z)]
    else:
        return z[np.argmin([c_m * _z + alpha * E([f_3t(_z - D_t_i, epoch + 1) for D_t_i in D_t], p_D_t) for _z in z])]


def S_1t(epoch):
    return np.inf


def S_2t(epoch):
    return y[np.argmin([g_1t(_y, epoch) for _y in y])]


def S_3t(epoch):
    return y[np.argmin([g_1t(_y, epoch) + g_2t(_y, epoch) for _y in y])]


def S_4t(epoch):
    return - np.inf


def f_t(v_0, v_1, v_2, v_3, epoch):
    return f_0t(v_0, epoch) + f_1t(v_1, epoch) + f_2t(v_2, epoch) + f_3t(v_3, epoch)
#%%
epoch = 1
max_epoch = 2
v_0 = 1
v_1 = 2
v_2 = 3
v_3 = 4
z = range(-(max_epoch + 1) * np.max(D_t), (max_epoch + 1) * np.max(K_t) + 1)
y = range(-(max_epoch + 1) * np.max(D_t), (max_epoch + 1) * np.max(K_t) + 1)
#%%
f_t(1, 2, 3, 4, epoch)
#%%
